// qotd.h    Header file for QOTD project
#ifndef WINSOCK2
#define WINSOCK2
#include <winsock2.h>
#pragma comment (lib, "ws2_32.lib")
#endif

const int  MAX_HOSTNAME  = 1024;
const int  MAX_PORT      = 128;
const int  MAX_RECV_BUF  = 2048;
const int  MAX_SEND_BUF  = 128;
const int  WAIT_TIME = 5;

int    wait(SOCKET, int, int);
SOCKET connectsock (char*, char*, char*);


#ifndef IOSTREAM
#define IOSTREAM
#include <iostream>
#endif

using namespace std;

SOCKET connectsock (char *host, char *service, char *protocol) {
	// Parameter List:
	// [IN] host    : String containing host as either a DNS name or an IP address
	// [IN] service : String containing service type as name (eg. ECHO) or port number (eg. 7)
	// [IN] protocol: String containing protocol type (either "tcp" or "udp")
	// Return value : SOCKET handle of connected socket.

	//	NOTES: 1. This function requires winsock.h header file and wsock32.lib library file.
	//         2. The host and service parameters are not used for UDP sockets.  Empty strings may be used.

	struct hostent  *phe;      /* pointer to host information entry     */
	struct servent  *pse;      /* pointer to service information entry  */
	struct sockaddr_in client; /* an Intenet endpoint address           */
	int				status;	   /* status of connection attempt          */
	SOCKET			s;		   /* socket handle                         */

	/* Allocate and connect socket - type depends upon protocol */
	if (_stricmp(protocol, "tcp") == 0) {
		/* Create a TCP socket */
		s = socket(AF_INET, SOCK_STREAM, 0);
		if (s < 0 || s == INVALID_SOCKET) {	
			std::cout << "Cannot create socket" << std::endl;
			return INVALID_SOCKET;
		}

		memset(&client, 0, sizeof(client));		/* zero the sockaddr_in struct */

		/* Set the address family to AF_INET (= Internet) */
		client.sin_family = AF_INET;

		/* Map service name (or number) to port number */
		short iport;     /* 16-bit port number */		
		iport = htons( (u_short) atoi(service));	/* 1st try to convert string to integer */
		if (iport == 0) {							/* if that doesn't work, call service function */
			pse = getservbyname(service, protocol);
			if (pse) {
				iport = pse->s_port;
			}
			else {
				std::cout << "Invalid service request." << std::endl;
				return INVALID_SOCKET;
			}
		}
		client.sin_port = iport;	/* Assign integer port # to field in client struct */


		/* Map host name to remote IP address, allowing for dotted quad */
		long ihost;	   /* 32-bit IP address number */
		ihost = inet_addr(host);	/* 1st try to convert from "dotted decimal notation" */
		if (ihost == INADDR_NONE) {	/* if that doesn't work, send request to DNS */
			phe = gethostbyname(host);
			if (phe)
				memmove(&ihost, phe->h_addr, phe->h_length);
			else {
				std::cout << "Invalid host" << std::endl;
				return INVALID_SOCKET;
			}
		}
		client.sin_addr.s_addr = ihost;	/* Assign integer IP address to field in client struct */


		/* Connect the socket */
		status = connect(s, (LPSOCKADDR) &client, sizeof(SOCKADDR));
		if (status == SOCKET_ERROR) {	
			std::cout << "Remote host/service not found - or connection refused" << std::endl;
			return INVALID_SOCKET;
		}
	}

	else if (_stricmp(protocol,"udp") == 0) {
		/* Create a UDP socket */
		s = socket(AF_INET, SOCK_DGRAM,  0);
		if (s < 0 || s == INVALID_SOCKET) {	
			std::cout << "Cannot create socket" << std::endl;
			return INVALID_SOCKET;
		}
	}

	else {	
		std::cout << "Invalid Protocol" << std::endl;
		return INVALID_SOCKET;
	}
		
	return s;
}


int wait( SOCKET s, int seconds, int msec ) 
{
	int    stat;					
	struct timeval stTimeOut;		
	fd_set stReadFDS;				
	fd_set stXcptFDS;				
	
	stTimeOut.tv_sec  = seconds;	
	stTimeOut.tv_usec = msec;		

	FD_ZERO( &stReadFDS );			
	FD_ZERO( &stXcptFDS );			
	FD_SET( s, &stReadFDS );			
	FD_SET( s, &stXcptFDS );			

	// Check to see if Read access is enabled for socket "s"
	stat = select( -1, &stReadFDS, NULL, &stXcptFDS, &stTimeOut );
	if ( stat == SOCKET_ERROR ) 
	{
		std::cout << std::endl << "wait() function failed" << std::endl;
		stat = 0;
	}
	else if ( stat > 0 ) 
	{
		if ( FD_ISSET( s, &stXcptFDS )) 
			stat = 0;
		
		if ( !FD_ISSET( s, &stReadFDS ))	
			stat = 0;
	}

	return stat;
}


void main (int argc, char** argv) {
	WORD wVersionRequired = 0x0101;
	WSADATA wsaData;
	int err = WSAStartup( wVersionRequired, &wsaData);
	if (err) {
		cout << "Unable to initialize Windows Socket library. " << endl;
		return;
	}

	char remote_host[MAX_HOSTNAME] = "";
	char remote_port[MAX_PORT]     = "";
	if (argc < 2) {		// Insufficient parameters on command line
		cout << "Usage: qotd_client [servername|ip_address] [port#|service]" << endl;
		return;
	}
	else if (argc == 2) {	// Get hostname from command line - assume standard port#
		strcpy_s(remote_host, MAX_HOSTNAME, argv[1]);
		strcpy_s(remote_port, MAX_PORT, "qotd");
	}
	else {				// Get hostname and port# (or service name) from command line
		strcpy_s(remote_host, MAX_HOSTNAME, argv[1]);
		strcpy_s(remote_port, MAX_PORT, argv[2]);
	}

	// Connect to remote server using TCP
	SOCKET socket = connectsock(remote_host, remote_port, "tcp");
	if (socket != INVALID_SOCKET) {
		char cmd[MAX_SEND_BUF];
		cout << "Command to send? ";
		cin.getline(cmd,MAX_SEND_BUF-1);
		send(socket, cmd, (int) strlen(cmd)+1, 0);
		int len = 1;
		while ( len>0 && wait(socket,WAIT_TIME,0)>0 ) {
			char recv_buf[MAX_RECV_BUF+1];
			len = recv(socket, recv_buf, MAX_RECV_BUF, 0);
			if (len > 0) {
				for (int i=0; i<len; i++) cout << recv_buf[i];
			}
		}
		closesocket(socket);
	}
	
	WSACleanup();
}